# -*- encoding: utf-8 -*-
"""
@File    :   noise.py
@Contact :   haoyu_deng@std.uestc.edu.cn


@Modify Time      @Author          @Version    @Desciption
------------      -------------    --------    -----------
2022/11/20 18:19   risingentropy      1.0         None
"""


class NoiseGeneratorBase:
    def __init__(self):
        """

        """

    def getNoise(self, shape: tuple, backend, device='cpu'):
        raise NotImplemented(
            "NoiseGeneratorBase is the super class of all noise generator, and is an abstract class that you cannot "
            "invoke getNoise method")


class GaussianNoiseGenerator(NoiseGeneratorBase):
    """

    """

    def __init__(self, mean=0, sigma=1):
        """

        :param mean:
        :param sigma:
        """
        super(GaussianNoiseGenerator, self).__init__()
        self.mean = mean
        self.sigma = sigma

    def getNoise(self, shape: tuple, backend, device='cpu'):
        """

        :param shape:
        :param backend:
        :return:
        """
        return backend.normal(loc=self.mean, scale=self.sigma, shape=shape, device=device)


class LaplaceNoiseGenerator(NoiseGeneratorBase):
    """

    """

    def __init__(self, loc=0, scale=1):
        """

        :param loc:
        :param scale:
        """
        super(LaplaceNoiseGenerator, self).__init__()
        self.loc = loc
        self.scale = scale

    def getNoise(self, shape: tuple, backend, device='cpu'):
        """


        :param shape:
        :param backend:
        :param device:
        :return:
        """
        return backend.laplace(loc=self.loc, scale=self.scale, shape=shape, device=device)


class MixedNoiseGenerator(NoiseGeneratorBase):
    def __init__(self, noiseGenerators: tuple, weights: tuple):
        """

        :param noiseGenerators:
        :param weights:
        """
        super().__init__()
        if len(noiseGenerators) != len(weights):
            raise ValueError("the length of weight doesn't match noiseGenerator!")
        self.noiseGenerators = noiseGenerators
        self.weights = list(weights)
        self.wei_sum = sum(weights)

        for i in range(0, len(self.weights)):
            if self.weights[i] < 0:
                raise ValueError("weight cannot be negative, the input weights are:" + str(weights))

            self.weights[i] = self.weights[i] / self.wei_sum
        self.prefix_sum = self.weights
        for i in range(1, len(self.weights)):
            self.prefix_sum[i] = self.prefix_sum[i - 1] + self.prefix_sum[i]

    def getNoise(self, shape: tuple, backend, device='cpu'):
        """


        :param shape:
        :param backend:
        :param device:
        :return:
        """
        result = backend.zeros(shape=shape)
        weight_map = backend.rand(shape)
        result = result + (weight_map < self.prefix_sum[0]) * self.noiseGenerators[0].getNoise(shape=shape,
                                                                                               backend='numpy',
                                                                                               device=device)

        for i in range(1, len(self.noiseGenerators)):
            result = result + backend.logical_and(weight_map >= self.prefix_sum[i - 1],
                                                  weight_map < self.prefix_sum[i]) * \
                     self.noiseGenerators[i].getNoise(shape=shape, backend='numpy', device=device)
        return result
